import math
from constants import BUY, SELL, OPEN, CLOSE, LONG
from vquant.strategy import Strategy


class CustomStrategy(Strategy):
    def __init__(self, cerebro, **kwargs):
        super(CustomStrategy, self).__init__(cerebro, **kwargs)

    @property
    def logger(self):
        return self.cerebro.logger

    def used_signal(self, signal):
        created, updated = self.signal.update_or_insert(dict(used=1), _id=dict(eq=signal._id))
        if updated:
            data = self.signal.query(_id=dict(eq=signal._id))
            data = data.iloc[0].to_dict()
        else:
            data = created.to_dict()
        self.cerebro.notify_signal(data)

    def stop_by_drawdown(self, position, history, tick):
        drawdown_rate = self.params.get('drawdown_rate')
        if drawdown_rate is not None:
            temp_df = history.loc[history.datetime.ge(position.created)]
            if position.direction == self.broker.position.Long:
                stop_price = temp_df['high'].max() * (1 - drawdown_rate / 100)
                return tick.close < stop_price
            if position.direction == self.broker.position.Short:
                stop_price = temp_df['low'].min() * (1 + drawdown_rate / 100)
                return tick.close > stop_price
        return False

    def stop_by_expired(self, position, history):
        position_expired_days = self.params.get('position_expired_days')
        if position_expired_days is not None:
            temp_df = history.loc[history.datetime.ge(position.created)]
            return temp_df.shape[0] > position_expired_days
        return False

    def stop_by_price(self, position, history, tick):
        stop_price_rate = self.params.get('stop_price_rate')
        if stop_price_rate is not None:
            temp_df = history.loc[history.datetime.lt(position.created)]
            if temp_df.shape[0] > 0:
                signal = temp_df.iloc[0]
                if position.direction == self.broker.position.Long:
                    stop_price = signal.high + signal.high * (stop_price_rate / 100)
                    return tick.close > stop_price
                else:
                    stop_price = signal.low - signal.low * (stop_price_rate / 100)
                    return tick.close < stop_price
        return False

    def stop_by_signal(self, position):
        if position.direction == self.broker.position.Long:
            signals = self.signal.query(
                symbol=dict(eq=position.symbol),
                side=dict(eq=SELL),
                flag=dict(eq=CLOSE),
                used=dict(eq=0)
            )
        else:
            signals = self.signal.query(
                symbol=dict(eq=position.symbol),
                side=dict(eq=BUY),
                flag=dict(eq=CLOSE),
                used=dict(eq=0)
            )
        for _, signal in signals.iterrows():
            self.used_signal(signal._id)
        return signals.shape[0] > 0

    def stop_by_target(self, position, history, tick):
        stop_rate = self.params.get('stop_rate')
        if stop_rate is not None:
            temp_df = history.loc[history.datetime.ge(position.created)]
            price = abs(position.cost / position.quantity / self.broker.quantity_multiple)
            target_rate = self.params.get('target_rate')
            slope_rate = self.params.get('slope_rate') or 0
            if position.direction == LONG:
                count = 0
                if target_rate:
                    upper = temp_df['high'].max()
                    count = (upper - price) / price * 100 / target_rate
                    count = math.floor(count) if count > 0 else 0
                stop_price = price * (1 - slope_rate / 100 + slope_rate / 100 * count)
                return tick.close <= stop_price
            else:
                count = 0
                if target_rate:
                    floor = temp_df['low'].min()
                    count = (floor - price) / price * 100 / target_rate
                    count = math.ceil(count) if count < 0 else 0
                stop_price = price * (1 + stop_rate / 100 + slope_rate / 100 * count)
                return tick.close >= stop_price
        return False

    def lock_position(self, position, history):
        lock_position_days = self.params.get('lock_position_days')
        if lock_position_days is not None:
            temp_df = history.loc[history.datetime.ge(position.created)]
            return temp_df.shape[0] < lock_position_days
        return False

    def lock_profit(self, position, tick):
        lock_profit_rate = self.params.get('lock_profit_rate')
        if lock_profit_rate is not None:
            profit = tick.open * position.quantity * self.broker.quantity_multiple - position.cost
            profit_rate = profit / abs(position.cost) * 100
            return profit_rate < lock_profit_rate
        return False

    def open_deviate(self, signal, tick):
        deviate = (tick.open - signal.d_close) / signal.d_close * 100
        min_open_deviate = self.params.get('min_open_deviate')
        if min_open_deviate is not None and deviate < min_open_deviate:
            self.logger.info(
                f"OpenDeviate: Symbol-{tick.symbol}, Deviate-{round(deviate, 2)}%, MinOpenDeviate-{min_open_deviate}%")
            return False
        max_open_deviate = self.params.get('max_open_deviate')
        if max_open_deviate is not None and deviate > max_open_deviate:
            self.logger.info(
                f"OpenDeviate: Symbol-{tick.symbol}, Deviate-{round(deviate, 2)}%, MaxOpenDeviate-{max_open_deviate}%")
            return False
        return True

    def price_deviate(self, signal, tick):
        deviate = (tick.close - signal.d_close) / signal.d_close * 100
        min_price_deviate = self.params.get('min_price_deviate')
        if min_price_deviate is not None and deviate < min_price_deviate:
            self.logger.info(
                f"PriceDeviate: Symbol-{tick.symbol}, Deviate-{round(deviate, 2)}%, MinPriceDeviate-{min_price_deviate}%")
            return False
        max_price_deviate = self.params.get('max_price_deviate')
        if max_price_deviate is not None and deviate > max_price_deviate:
            self.logger.info(
                f"PriceDeviate: Symbol-{tick.symbol}, Deviate-{round(deviate, 2)}%, MaxPriceDeviate-{max_price_deviate}%")
            return False
        return True

    def close_position(self, tick, position, message):
        side = SELL if position.direction == LONG else BUY
        self.broker.create_order(tick.symbol, position.direction, side, tick.price, position.quantity)
        self.logger.info(
            f"ExecuteClosingPosition: Symbol-{tick.symbol}, Price-{tick.price}, Quantity-{position.quantity}, Message: {message}")

    def exec_close(self, positions, ticks):
        positions = positions.loc[positions.symbol.isin(ticks.symbol.values)]
        for _, position in positions.iterrows():
            tick = ticks.loc[position.symbol]
            history = self.datafeed.get_history(position.symbol)
            if self.lock_position(position, history):
                continue
            if self.lock_profit(position, tick):
                continue
            if self.stop_by_expired(position, history):
                self.close_position(tick, position, 'StopByExpired')
                continue
            if self.stop_by_signal(position):
                self.close_position(tick, position, 'StopBySignal')
                continue
            if self.stop_by_drawdown(position, history, tick):
                self.close_position(tick, position, 'StopByDrawdown')
                continue
            if self.stop_by_price(position, history, tick):
                self.close_position(tick, position, 'StopByPrice')
                continue
            if self.stop_by_target(position, history, tick):
                self.close_position(tick, position, 'StopByTarget')
                continue

    def exec_open(self, positions, ticks):
        if not self.signal.dataframe.shape[0]:
            return
        signals = self.signal.query(
            datetime=dict(eq=self.cerebro.datafeed.get_previous_trading_date()),
            symbol=dict(isin=ticks.symbol.values),
            flag=dict(eq=OPEN),
            used=dict(eq=0)
        )
        signals = signals.loc[~signals.symbol.isin(positions.symbol.values)]
        entries = int(self.params['entry']) - positions.shape[0]
        available = self.broker.asset_info["available"] - self.broker.asset_info["frozen"]
        counted = math.floor(available / float(self.params['amount']))
        notional = float(self.params['amount'])
        if counted < entries:
            notional = available / entries
            counted = entries
        for _, signal in signals.iterrows():
            self.logger.info(f"OpeningAvailable: Funds-{available}, Positions-{counted}")
            if not counted:
                break
            tick = ticks.loc[signal.symbol]
            if not self.open_deviate(signal, tick):
                self.used_signal(signal)
                continue
            if not self.price_deviate(signal, tick):
                self.used_signal(signal)
                continue
            quantity = math.floor(notional / tick.close)
            if quantity > 100:
                quantity = math.floor(quantity / 100) * 100
            else:
                self.logger.info(f"AvailableForPurchase: Symbol-{tick.symbol}, Quantity-{quantity}")
                continue
            price = self.broker.create_order(tick, BUY, quantity)
            self.logger.info(f"ExecuteOpeningSignal: Symbol-{tick.symbol}, Price-{price}, Quantity-{quantity}")
            self.used_signal(signal)
            counted -= 1

    def on_next(self, ticks):
        positions = self.broker.position.query(quantity=dict(ne=0))
        self.exec_close(positions, ticks)
        self.exec_open(positions, ticks)
